import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp
from scipy.linalg import solve_continuous_are

class WheeledLegRobot:
    def __init__(self):
        # Robot parameters
        self.m_body = 2.0      # body mass kg
        self.m_leg = 0.5       # leg mass kg
        self.l_hip = 0.3       # hip joint length m
        self.l_knee = 0.3      # knee joint length m
        self.I_body = 0.1      # body moment of inertia
        self.g = 9.81          # gravity
        
        # Leg parameters - 4 legs with suspension
        self.num_legs = 4
        self.leg_positions = np.array([
            [0.2, 0.15, 0],   # front left
            [0.2, -0.15, 0],  # front right  
            [-0.2, 0.15, 0],  # rear left
            [-0.2, -0.15, 0]  # rear right
        ])
        
        # Leg state: [leg1_height, leg2_height, leg3_height, leg4_height, ... velocities]
        self.leg_state_dim = 8  # 4 positions + 4 velocities
        self.leg_control_dim = 4
        
        # Main state: [x, y, z, theta, phi, psi, x_dot, y_dot, z_dot, theta_dot, phi_dot, psi_dot]
        self.state_dim = 12
        self.control_dim = 6 + 4  # 6 body controls + 4 leg controls
        
    def dynamics(self, state, control_torques, t, disturbance_type='none'):
        """Robot dynamics equations with leg suspension and disturbances"""
        # Extract main body state
        x, y, z, theta, phi, psi, x_dot, y_dot, z_dot, theta_dot, phi_dot, psi_dot = state[0:12]
        
        # Extract leg states (if provided)
        if len(state) > 12:
            leg_heights = state[12:16]
            leg_velocities = state[16:20]
        else:
            leg_heights = np.array([0.3, 0.3, 0.3, 0.3])  # default leg heights
            leg_velocities = np.zeros(4)
        
        # Separate body and leg controls
        body_control = control_torques[0:6]
        leg_control = control_torques[6:10]
        
        # Simplified dynamics model for body
        M_body = np.diag([self.m_body, self.m_body, self.m_body, 
                         self.I_body, self.I_body, self.I_body])
        
        # Coriolis and gravity terms
        C_body = np.array([0, 0, self.m_body * self.g, 0, 0, 0])
        
        # Leg suspension dynamics - 修复参数设置
        leg_stiffness = 100.0  # N/m
        leg_damping = 10.0     # Ns/m
        leg_mass = 0.5         # kg per leg - 增加腿的质量避免数值问题
        
        # Ground contact forces
        ground_forces = self.calculate_ground_forces(z, leg_heights, leg_velocities)
        
        # Leg accelerations - 添加数值限制
        leg_forces = leg_control - leg_stiffness * leg_heights - leg_damping * leg_velocities + ground_forces
        leg_forces = np.clip(leg_forces, -1000, 1000)  # 限制力的大小
        leg_accelerations = leg_forces / leg_mass
        
        # Body acceleration with leg suspension effects
        total_leg_force = np.sum(ground_forces)
        body_disturbance_z = -total_leg_force / self.m_body
        
        # 限制body_disturbance_z的大小
        body_disturbance_z = np.clip(body_disturbance_z, -100, 100)
        
        # Add disturbances based on type
        disturbance = self.get_disturbance(t, disturbance_type)
        disturbance[2] += body_disturbance_z  # Add leg suspension effect to z disturbance
        
        # Control input mapping
        B = np.eye(6)
        
        # Acceleration calculation with disturbance
        try:
            accelerations = np.linalg.solve(M_body, B @ body_control - C_body + disturbance)
        except np.linalg.LinAlgError:
            # 如果出现数值问题，使用伪逆
            accelerations = np.linalg.pinv(M_body) @ (B @ body_control - C_body + disturbance)
        
        # 限制加速度的大小
        accelerations = np.clip(accelerations, -100, 100)
        
        # State derivatives
        state_dot = np.zeros(len(state))
        state_dot[0:6] = state[6:12]  # position derivatives = velocities
        state_dot[6:12] = accelerations  # velocity derivatives = accelerations
        
        # Leg state derivatives
        if len(state) > 12:
            state_dot[12:16] = leg_velocities
            state_dot[16:20] = leg_accelerations
        
        return state_dot
    
    def calculate_ground_forces(self, body_height, leg_heights, leg_velocities):
        """Calculate ground reaction forces for each leg"""
        ground_forces = np.zeros(4)
        ground_stiffness = 1000.0  # N/m for ground contact
        ground_damping = 50.0      # Ns/m
        
        for i in range(4):
            # Total height from ground to body through leg
            total_height = body_height + leg_heights[i]
            
            # Ground penetration depth
            penetration = max(0, -total_height)  # Positive when penetrating ground
            
            if penetration > 0:
                # Ground reaction force (spring-damper model)
                ground_force = (ground_stiffness * penetration - 
                              ground_damping * leg_velocities[i])
                ground_forces[i] = max(0, ground_force)
            else:
                ground_forces[i] = 0
                
        return ground_forces
    
    def get_leg_kinematics(self, state):
        """Calculate leg positions and orientations"""
        if len(state) <= 12:
            return None, None, None
            
        x, y, z, theta, phi, psi = state[0:6]
        leg_heights = state[12:16]
        
        # Base positions in body frame
        base_positions = self.leg_positions.copy()
        
        # Apply body orientation
        rot_x = np.array([
            [1, 0, 0],
            [0, np.cos(phi), -np.sin(phi)],
            [0, np.sin(phi), np.cos(phi)]
        ])
        
        rot_y = np.array([
            [np.cos(theta), 0, np.sin(theta)],
            [0, 1, 0],
            [-np.sin(theta), 0, np.cos(theta)]
        ])
        
        rot_z = np.array([
            [np.cos(psi), -np.sin(psi), 0],
            [np.sin(psi), np.cos(psi), 0],
            [0, 0, 1]
        ])
        
        rotation_matrix = rot_z @ rot_y @ rot_x
        
        # Transform to world frame
        world_positions = []
        foot_positions = []
        
        for i in range(4):
            base_world = rotation_matrix @ base_positions[i] + np.array([x, y, z])
            foot_world = base_world - np.array([0, 0, leg_heights[i]])
            
            world_positions.append(base_world)
            foot_positions.append(foot_world)
        
        return np.array(world_positions), np.array(foot_positions), leg_heights
    
    def get_disturbance(self, t, disturbance_type):
        """Generate different types of disturbances"""
        disturbance = np.zeros(6)
        
        if disturbance_type == 'none':
            return disturbance
            
        elif disturbance_type == 'constant':
            # Constant disturbance force - 减小幅度
            disturbance = np.array([1.0, 0.5, 1.5, 0.2, 0.2, 0.2])
            
        elif disturbance_type == 'sinusoidal':
            # Sinusoidal disturbance - 减小幅度
            amplitude = 1.5
            frequency = 1.0
            disturbance = amplitude * np.array([
                np.sin(2*np.pi*frequency*t),
                np.cos(2*np.pi*frequency*t),
                0.3*np.sin(2*np.pi*frequency*t + np.pi/4),
                0.1*np.cos(2*np.pi*frequency*t),
                0.1*np.sin(2*np.pi*frequency*t),
                0.05*np.cos(2*np.pi*frequency*t)
            ])
            
        elif disturbance_type == 'impulse':
            # Impulse disturbances at specific times - 减小幅度
            if 2.0 <= t < 2.1:
                disturbance = np.array([5, 4, 3, 1, 1, 1])
            elif 5.0 <= t < 5.1:
                disturbance = np.array([-4, 5, 2, 0.5, 0.5, 0.5])
            else:
                disturbance = np.zeros(6)
                
        elif disturbance_type == 'random':
            # Random noise disturbance - 减小幅度
            disturbance = 1.0 * np.random.randn(6)
            
        elif disturbance_type == 'step':
            # Step disturbance - 减小幅度
            if t >= 3.0:
                disturbance = np.array([2.0, 1.5, 1.0, 0.5, 0.5, 0.5])
                
        elif disturbance_type == 'complex':
            # Complex combined disturbance - 减小幅度
            if t < 3.0:
                # Constant
                disturbance = np.array([0.5, 0.3, 1.0, 0.15, 0.15, 0.15])
            elif 3.0 <= t < 6.0:
                # Sinusoidal
                amplitude = 1.0
                disturbance = amplitude * np.array([
                    np.sin(2*np.pi*1.0*t),
                    np.cos(2*np.pi*1.0*t),
                    0,
                    0.1*np.sin(2*np.pi*1.0*t),
                    0.1*np.cos(2*np.pi*1.0*t),
                    0
                ])
            else:
                # Random
                disturbance = 0.8 * np.random.randn(6)
                
        return disturbance

class PIDController:
    def __init__(self, kp, ki, kd, dt, include_legs=False):
        self.kp = kp
        self.ki = ki
        self.kd = kd
        self.dt = dt
        self.include_legs = include_legs
        
        control_dim = 10 if include_legs else 6
        self.integral = np.zeros(control_dim)
        self.previous_error = np.zeros(control_dim)
    
    def compute_control(self, desired_state, current_state):
        if self.include_legs:
            # For body + leg control
            error = desired_state[0:10] - current_state[0:10]
        else:
            # For body control only
            error = desired_state[0:6] - current_state[0:6]
            
        error_derivative = (error - self.previous_error) / self.dt
        
        self.integral += error * self.dt
        self.integral = np.clip(self.integral, -50, 50)  # Anti-windup
        
        control = (self.kp * error + 
                  self.ki * self.integral + 
                  self.kd * error_derivative)
        
        self.previous_error = error.copy()
        
        if not self.include_legs:
            # Add zero leg controls
            control = np.concatenate([control, np.zeros(4)])
            
        return control

class SMCController:
    def __init__(self, lambda_val, k, phi, dt, include_legs=False):
        self.lambda_val = lambda_val
        self.k = k
        self.phi = phi
        self.dt = dt
        self.include_legs = include_legs
        
        control_dim = 10 if include_legs else 6
        self.previous_error = np.zeros(control_dim)
    
    def compute_control(self, desired_state, current_state):
        if self.include_legs:
            error = desired_state[0:10] - current_state[0:10]
        else:
            error = desired_state[0:6] - current_state[0:6]
            
        error_derivative = (error - self.previous_error) / self.dt
        
        # Sliding surface
        s = error_derivative + self.lambda_val * error
        
        # Control law
        sat_s = np.tanh(s * 5)
        control = self.k * sat_s + self.phi * s
        
        self.previous_error = error.copy()
        
        if not self.include_legs:
            control = np.concatenate([control, np.zeros(4)])
            
        return control

class LQRController:
    def __init__(self, A, B, Q, R, include_legs=False):
        self.include_legs = include_legs
        if not include_legs:
            # Use only body states for LQR
            P = solve_continuous_are(A, B, Q, R)
            self.K = np.linalg.solve(R, B.T @ P)
        else:
            self.K = np.zeros((10, A.shape[0]))  # Placeholder
    
    def compute_control(self, desired_state, current_state):
        if self.include_legs:
            error = current_state - desired_state
            control = -self.K @ error
        else:
            error = current_state[0:12] - desired_state[0:12]
            control_body = -self.K @ error
            control = np.concatenate([control_body, np.zeros(4)])
            
        return control

def simulate_controller_with_disturbance(controller, robot, desired_trajectory, t_span, dt, disturbance_type):
    """Simulate specific controller with disturbance and leg suspension"""
    t_eval = np.arange(t_span[0], t_span[1], dt)
    n_steps = len(t_eval)
    
    # Initialize state records - now including leg states
    full_state_dim = robot.state_dim + robot.leg_state_dim
    states = np.zeros((n_steps, full_state_dim))
    controls = np.zeros((n_steps, robot.control_dim))
    errors = np.zeros((n_steps, 10))  # 6 body + 4 leg errors
    disturbances = np.zeros((n_steps, 6))
    leg_heights_history = np.zeros((n_steps, 4))
    
    # Initial state with legs extended
    state = np.zeros(full_state_dim)
    state[0:12] = np.array([0, 0, 0.5, 0, 0, 0, 0, 0, 0, 0, 0, 0])  # body state
    state[12:16] = np.array([0.3, 0.3, 0.3, 0.3])  # leg heights
    state[16:20] = np.zeros(4)  # leg velocities
    
    for i, t in enumerate(t_eval):
        # Get desired state (including leg heights if needed)
        desired_full_state = desired_trajectory(t)
        
        # Calculate control input
        control = controller.compute_control(desired_full_state, state)
        
        # Limit control input - 减小控制限制
        control = np.clip(control, -15, 15)
        
        # Get current disturbance
        current_disturbance = robot.get_disturbance(t, disturbance_type)
        
        # Record state and control
        states[i] = state
        controls[i] = control
        errors[i, 0:6] = desired_full_state[0:6] - state[0:6]
        if len(desired_full_state) > 10:
            errors[i, 6:10] = desired_full_state[6:10] - state[12:16]
        disturbances[i] = current_disturbance
        leg_heights_history[i] = state[12:16]
        
        # Integrate dynamics with disturbance
        state_dot = robot.dynamics(state, control, t, disturbance_type)
        state = state + state_dot * dt
        
        # 添加状态限制，防止数值爆炸
        state = np.clip(state, -1e6, 1e6)
    
    return t_eval, states, controls, errors, disturbances, leg_heights_history

def desired_figure8_trajectory_with_legs(t):
    """Figure-8 trajectory with adaptive leg heights"""
    a = 1.0
    b = 0.7
    omega = 0.4
    
    # Body trajectory
    x = a * np.sin(omega * t)
    y = b * np.sin(2 * omega * t)
    z = 0.5 + 0.1 * np.sin(omega * t + np.pi/4)
    
    # Orientation
    theta = np.arctan2(y, x)
    phi = 0.1 * np.sin(omega * t)
    psi = 0.05 * np.cos(omega * t)
    
    # Velocities
    x_dot = a * omega * np.cos(omega * t)
    y_dot = 2 * b * omega * np.cos(2 * omega * t)
    z_dot = 0.1 * omega * np.cos(omega * t + np.pi/4)
    
    # Angular velocities
    theta_dot, phi_dot, psi_dot = 0, 0.1 * omega * np.cos(omega * t), -0.05 * omega * np.sin(omega * t)
    
    # Adaptive leg heights based on terrain and motion
    base_leg_height = 0.3
    terrain_adaptation = 0.05 * np.sin(2 * omega * t)  # Simulate uneven terrain
    motion_adaptation = 0.02 * np.sin(4 * omega * t)   # Leg cycling for walking
    
    leg_heights = base_leg_height + terrain_adaptation + motion_adaptation * np.array([1, -1, -1, 1])
    
    # Leg velocities (simplified)
    leg_velocities = np.zeros(4)
    
    return np.concatenate([
        [x, y, z, theta, phi, psi, x_dot, y_dot, z_dot, theta_dot, phi_dot, psi_dot],
        leg_heights, leg_velocities
    ])

def desired_rough_terrain_trajectory(t):
    """Trajectory with rough terrain requiring leg adaptation"""
    # Body follows simple path
    x = 0.8 * t
    y = 0.2 * np.sin(t)
    z = 0.5  # Body tries to maintain height
    
    # Orientation
    theta = 0
    phi = 0.05 * np.sin(2 * t)
    psi = 0
    
    # Velocities
    x_dot = 0.8
    y_dot = 0.2 * np.cos(t)
    z_dot = 0
    
    # Angular velocities
    theta_dot, phi_dot, psi_dot = 0, 0.1 * np.cos(2 * t), 0
    
    # Complex leg motion for rough terrain
    base_height = 0.3
    # Simulate stepping over obstacles
    front_left_height = base_height + 0.1 * np.sin(4 * t)
    front_right_height = base_height + 0.1 * np.sin(4 * t + np.pi)
    rear_left_height = base_height + 0.1 * np.sin(4 * t + np.pi/2)
    rear_right_height = base_height + 0.1 * np.sin(4 * t + 3*np.pi/2)
    
    leg_heights = np.array([front_left_height, front_right_height, rear_left_height, rear_right_height])
    leg_velocities = np.zeros(4)
    
    return np.concatenate([
        [x, y, z, theta, phi, psi, x_dot, y_dot, z_dot, theta_dot, phi_dot, psi_dot],
        leg_heights, leg_velocities
    ])

def plot_robustness_comparison_with_legs(t_pid, t_smc, t_lqr, states_pid, states_smc, states_lqr,
                                       errors_pid, errors_smc, errors_lqr, disturbances, 
                                       leg_heights_pid, leg_heights_smc, leg_heights_lqr,
                                       disturbance_type, trajectory_type='figure8'):
    """Plot robustness comparison results including leg information"""
    
    fig, axes = plt.subplots(3, 4, figsize=(20, 14))
    fig.suptitle(f'Robustness Test: {disturbance_type.upper()} Disturbance - {trajectory_type.upper()} Trajectory', 
                 fontsize=16, fontweight='bold')
    
    # Generate reference trajectory
    t_ref = np.linspace(0, 10, 500)
    if trajectory_type == 'figure8':
        x_ref = 1.0 * np.sin(0.4 * t_ref)
        y_ref = 0.7 * np.sin(0.8 * t_ref)
        z_ref = 0.5 + 0.1 * np.sin(0.4 * t_ref + np.pi/4)
    else:
        x_ref = 0.8 * t_ref
        y_ref = 0.2 * np.sin(t_ref)
        z_ref = 0.5 * np.ones_like(t_ref)
    
    # Position tracking
    axes[0,0].plot(t_pid, states_pid[:, 0], 'r-', label='PID', linewidth=2)
    axes[0,0].plot(t_smc, states_smc[:, 0], 'g-', label='SMC', linewidth=2)
    axes[0,0].plot(t_lqr, states_lqr[:, 0], 'b-', label='LQR', linewidth=2)
    axes[0,0].plot(t_ref, x_ref, 'k--', label='Desired', linewidth=1, alpha=0.7)
    axes[0,0].set_ylabel('X Position (m)', fontsize=12, labelpad=10)
    axes[0,0].set_title('X-Axis Position Tracking', fontsize=12, pad=15)
    axes[0,0].legend(fontsize=10)
    axes[0,0].grid(True)
    
    # Y position
    axes[0,1].plot(t_pid, states_pid[:, 1], 'r-', label='PID', linewidth=2)
    axes[0,1].plot(t_smc, states_smc[:, 1], 'g-', label='SMC', linewidth=2)
    axes[0,1].plot(t_lqr, states_lqr[:, 1], 'b-', label='LQR', linewidth=2)
    axes[0,1].plot(t_ref, y_ref, 'k--', label='Desired', linewidth=1, alpha=0.7)
    axes[0,1].set_ylabel('Y Position (m)', fontsize=12, labelpad=10)
    axes[0,1].set_title('Y-Axis Position Tracking', fontsize=12, pad=15)
    axes[0,1].legend(fontsize=10)
    axes[0,1].grid(True)
    
    # Z position (body height)
    axes[0,2].plot(t_pid, states_pid[:, 2], 'r-', label='PID', linewidth=2)
    axes[0,2].plot(t_smc, states_smc[:, 2], 'g-', label='SMC', linewidth=2)
    axes[0,2].plot(t_lqr, states_lqr[:, 2], 'b-', label='LQR', linewidth=2)
    axes[0,2].plot(t_ref, z_ref, 'k--', label='Desired', linewidth=1, alpha=0.7)
    axes[0,2].set_ylabel('Body Height (m)', fontsize=12, labelpad=10)
    axes[0,2].set_title('Body Height Tracking', fontsize=12, pad=15)
    axes[0,2].legend(fontsize=10)
    axes[0,2].grid(True)
    
    # Leg heights
    axes[0,3].plot(t_pid, leg_heights_pid[:, 0], 'r-', label='PID FL', linewidth=1.5, alpha=0.7)
    axes[0,3].plot(t_pid, leg_heights_pid[:, 1], 'r--', label='PID FR', linewidth=1.5, alpha=0.7)
    axes[0,3].plot(t_smc, leg_heights_smc[:, 0], 'g-', label='SMC FL', linewidth=1.5, alpha=0.7)
    axes[0,3].plot(t_smc, leg_heights_smc[:, 1], 'g--', label='SMC FR', linewidth=1.5, alpha=0.7)
    axes[0,3].plot(t_lqr, leg_heights_lqr[:, 0], 'b-', label='LQR FL', linewidth=1.5, alpha=0.7)
    axes[0,3].plot(t_lqr, leg_heights_lqr[:, 1], 'b--', label='LQR FR', linewidth=1.5, alpha=0.7)
    axes[0,3].set_ylabel('Front Leg Heights (m)', fontsize=12, labelpad=10)
    axes[0,3].set_title('Front Leg Adaptation', fontsize=12, pad=15)
    axes[0,3].legend(fontsize=8, ncol=2)
    axes[0,3].grid(True)
    
    # Position Error Norm
    error_norm_pid = np.linalg.norm(errors_pid[:, 0:3], axis=1)
    error_norm_smc = np.linalg.norm(errors_smc[:, 0:3], axis=1)
    error_norm_lqr = np.linalg.norm(errors_lqr[:, 0:3], axis=1)
    
    axes[1,0].plot(t_pid, error_norm_pid, 'r-', label='PID', linewidth=2)
    axes[1,0].plot(t_smc, error_norm_smc, 'g-', label='SMC', linewidth=2)
    axes[1,0].plot(t_lqr, error_norm_lqr, 'b-', label='LQR', linewidth=2)
    axes[1,0].set_ylabel('Position Error Norm (m)', fontsize=12, labelpad=10)
    axes[1,0].set_xlabel('Time (s)', fontsize=12, labelpad=10)
    axes[1,0].set_title('Body Position Error', fontsize=12, pad=15)
    axes[1,0].legend(fontsize=10)
    axes[1,0].grid(True)
    
    # Leg height errors
    leg_error_norm_pid = np.linalg.norm(errors_pid[:, 6:10], axis=1)
    leg_error_norm_smc = np.linalg.norm(errors_smc[:, 6:10], axis=1)
    leg_error_norm_lqr = np.linalg.norm(errors_lqr[:, 6:10], axis=1)
    
    axes[1,1].plot(t_pid, leg_error_norm_pid, 'r-', label='PID', linewidth=2)
    axes[1,1].plot(t_smc, leg_error_norm_smc, 'g-', label='SMC', linewidth=2)
    axes[1,1].plot(t_lqr, leg_error_norm_lqr, 'b-', label='LQR', linewidth=2)
    axes[1,1].set_ylabel('Leg Height Error Norm (m)', fontsize=12, labelpad=10)
    axes[1,1].set_xlabel('Time (s)', fontsize=12, labelpad=10)
    axes[1,1].set_title('Leg Height Error', fontsize=12, pad=15)
    axes[1,1].legend(fontsize=10)
    axes[1,1].grid(True)
    
    # XY trajectory
    axes[1,2].plot(states_pid[:, 0], states_pid[:, 1], 'r-', label='PID', linewidth=2, alpha=0.7)
    axes[1,2].plot(states_smc[:, 0], states_smc[:, 1], 'g-', label='SMC', linewidth=2, alpha=0.7)
    axes[1,2].plot(states_lqr[:, 0], states_lqr[:, 1], 'b-', label='LQR', linewidth=2, alpha=0.7)
    axes[1,2].plot(x_ref, y_ref, 'k--', label='Desired', linewidth=2)
    axes[1,2].set_xlabel('X (m)', fontsize=12, labelpad=10)
    axes[1,2].set_ylabel('Y (m)', fontsize=12, labelpad=10)
    axes[1,2].set_title('XY Trajectory Tracking', fontsize=12, pad=15)
    axes[1,2].legend(fontsize=10)
    axes[1,2].grid(True)
    axes[1,2].axis('equal')
    
    # Disturbance profile
    disturbance_norm = np.linalg.norm(disturbances, axis=1)
    axes[1,3].plot(t_pid, disturbance_norm, 'k-', linewidth=2)
    axes[1,3].set_ylabel('Disturbance Magnitude', fontsize=12, labelpad=10)
    axes[1,3].set_xlabel('Time (s)', fontsize=12, labelpad=10)
    axes[1,3].set_title('Applied Disturbance', fontsize=12, pad=15)
    axes[1,3].grid(True)
    
    # Rear leg heights
    axes[2,0].plot(t_pid, leg_heights_pid[:, 2], 'r-', label='PID RL', linewidth=1.5, alpha=0.7)
    axes[2,0].plot(t_pid, leg_heights_pid[:, 3], 'r--', label='PID RR', linewidth=1.5, alpha=0.7)
    axes[2,0].plot(t_smc, leg_heights_smc[:, 2], 'g-', label='SMC RL', linewidth=1.5, alpha=0.7)
    axes[2,0].plot(t_smc, leg_heights_smc[:, 3], 'g--', label='SMC RR', linewidth=1.5, alpha=0.7)
    axes[2,0].plot(t_lqr, leg_heights_lqr[:, 2], 'b-', label='LQR RL', linewidth=1.5, alpha=0.7)
    axes[2,0].plot(t_lqr, leg_heights_lqr[:, 3], 'b--', label='LQR RR', linewidth=1.5, alpha=0.7)
    axes[2,0].set_ylabel('Rear Leg Heights (m)', fontsize=12, labelpad=10)
    axes[2,0].set_xlabel('Time (s)', fontsize=12, labelpad=10)
    axes[2,0].set_title('Rear Leg Adaptation', fontsize=12, pad=15)
    axes[2,0].legend(fontsize=8, ncol=2)
    axes[2,0].grid(True)
    
    # Body orientation
    axes[2,1].plot(t_pid, np.degrees(states_pid[:, 3]), 'r-', label='PID Roll', linewidth=1.5)
    axes[2,1].plot(t_smc, np.degrees(states_smc[:, 3]), 'g-', label='SMC Roll', linewidth=1.5)
    axes[2,1].plot(t_lqr, np.degrees(states_lqr[:, 3]), 'b-', label='LQR Roll', linewidth=1.5)
    axes[2,1].set_ylabel('Body Roll (deg)', fontsize=12, labelpad=10)
    axes[2,1].set_xlabel('Time (s)', fontsize=12, labelpad=10)
    axes[2,1].set_title('Body Orientation', fontsize=12, pad=15)
    axes[2,1].legend(fontsize=10)
    axes[2,1].grid(True)
    
    # 3D visualization of robot posture at specific time
    time_sample = len(t_pid) // 2
    sample_states = [states_pid[time_sample], states_smc[time_sample], states_lqr[time_sample]]
    colors = ['red', 'green', 'blue']
    labels = ['PID', 'SMC', 'LQR']
    
    ax_3d = axes[2,2]
    ax_3d.remove()
    ax_3d = fig.add_subplot(3, 4, 11, projection='3d')
    
    robot = WheeledLegRobot()
    for i, state in enumerate(sample_states):
        base_pos, foot_pos, leg_heights = robot.get_leg_kinematics(state)
        if base_pos is not None:
            # Plot body
            ax_3d.scatter(state[0], state[1], state[2], color=colors[i], s=100, label=labels[i])
            # Plot legs
            for j in range(4):
                ax_3d.plot([base_pos[j, 0], foot_pos[j, 0]], 
                          [base_pos[j, 1], foot_pos[j, 1]], 
                          [base_pos[j, 2], foot_pos[j, 2]], 
                          color=colors[i], linewidth=2, alpha=0.7)
    
    ax_3d.set_xlabel('X (m)')
    ax_3d.set_ylabel('Y (m)')
    ax_3d.set_zlabel('Z (m)')
    ax_3d.set_title('Robot Posture (Sample)', fontsize=12, pad=15)
    ax_3d.legend()
    
    # Performance metrics
    metrics = calculate_robustness_metrics(errors_pid, errors_smc, errors_lqr, disturbance_type)
    robustness_margins = metrics['Robustness Margin']
    controllers = ['PID', 'SMC', 'LQR']
    
    axes[2,3].bar(controllers, robustness_margins, color=['red', 'green', 'blue'], alpha=0.7)
    axes[2,3].set_ylabel('Robustness Margin', fontsize=12, labelpad=10)
    axes[2,3].set_title('Controller Robustness', fontsize=12, pad=15)
    axes[2,3].grid(True, alpha=0.3)
    axes[2,3].set_ylim(0, 1)
    
    for i, margin in enumerate(robustness_margins):
        axes[2,3].text(i, margin + 0.02, f'{margin:.3f}', ha='center', va='bottom', fontsize=11)
    
    plt.tight_layout(pad=3.0)
    plt.subplots_adjust(top=0.93)
    plt.show()
    
    return metrics

def run_robustness_test_with_legs(disturbance_type, trajectory_type='figure8'):
    """Run robustness test with leg suspension"""
    print(f"\n{'='*60}")
    print(f"ROBUSTNESS TEST WITH LEG SUSPENSION: {disturbance_type.upper()} DISTURBANCE")
    print(f"TRAJECTORY: {trajectory_type.upper()}")
    print(f"{'='*60}")
    
    # Initialize robot with leg suspension
    robot = WheeledLegRobot()
    dt = 0.01
    t_span = (0, 10)
    
    # Select trajectory
    if trajectory_type == 'figure8':
        desired_trajectory_func = desired_figure8_trajectory_with_legs
    elif trajectory_type == 'rough':
        desired_trajectory_func = desired_rough_terrain_trajectory
    else:
        desired_trajectory_func = desired_figure8_trajectory_with_legs
    
    # Controllers with leg control - 调整控制器参数
    pid = PIDController(
        kp=np.array([15, 15, 20, 10, 10, 10, 6, 6, 6, 6]),  # 减小增益
        ki=np.array([0.1, 0.1, 0.5, 0.1, 0.1, 0.1, 0.05, 0.05, 0.05, 0.05]),  # 减小积分增益
        kd=np.array([4, 4, 8, 3, 3, 3, 1.5, 1.5, 1.5, 1.5]),  # 减小微分增益
        dt=dt,
        include_legs=True
    )
    
    smc = SMCController(
        lambda_val=2.0,  # 减小lambda
        k=np.array([8, 8, 12, 6, 6, 6, 4, 4, 4, 4]),  # 减小增益
        phi=np.array([2, 2, 3, 1.5, 1.5, 1.5, 0.8, 0.8, 0.8, 0.8]),  # 减小增益
        dt=dt,
        include_legs=True
    )
    
    # LQR (body only for simplicity)
    A = np.zeros((12, 12))
    A[0:6, 6:12] = np.eye(6)
    A[8, 8] = -0.2
    
    B = np.zeros((12, 6))
    B[6:12, 0:6] = np.eye(6) / robot.m_body
    B[9:12, 3:6] = np.eye(3) / robot.I_body
    
    Q = np.diag([10, 10, 20, 6, 6, 6, 1.5, 1.5, 2, 0.8, 0.8, 0.8])  # 减小Q矩阵权重
    R = np.diag([0.1, 0.1, 0.1, 0.1, 0.1, 0.1])  # 增加R矩阵权重
    
    lqr = LQRController(A, B, Q, R, include_legs=False)
    
    # Run simulations
    print("Running simulations with leg suspension...")
    t_pid, states_pid, controls_pid, errors_pid, dist_pid, legs_pid = simulate_controller_with_disturbance(
        pid, robot, desired_trajectory_func, t_span, dt, disturbance_type)
    
    t_smc, states_smc, controls_smc, errors_smc, dist_smc, legs_smc = simulate_controller_with_disturbance(
        smc, robot, desired_trajectory_func, t_span, dt, disturbance_type)
    
    t_lqr, states_lqr, controls_lqr, errors_lqr, dist_lqr, legs_lqr = simulate_controller_with_disturbance(
        lqr, robot, desired_trajectory_func, t_span, dt, disturbance_type)
    
    # Plot results
    metrics = plot_robustness_comparison_with_legs(t_pid, t_smc, t_lqr, 
                                                 states_pid, states_smc, states_lqr,
                                                 errors_pid, errors_smc, errors_lqr,
                                                 dist_pid, legs_pid, legs_smc, legs_lqr,
                                                 disturbance_type, trajectory_type)
    
    return metrics

def calculate_robustness_metrics(errors_pid, errors_smc, errors_lqr, disturbance_type):
    """Calculate robustness performance metrics"""
    
    error_norm_pid = np.linalg.norm(errors_pid[:, 0:3], axis=1)
    error_norm_smc = np.linalg.norm(errors_smc[:, 0:3], axis=1)
    error_norm_lqr = np.linalg.norm(errors_lqr[:, 0:3], axis=1)
    
    # Individual axis errors
    x_error_pid = np.abs(errors_pid[:, 0])
    y_error_pid = np.abs(errors_pid[:, 1])
    z_error_pid = np.abs(errors_pid[:, 2])
    
    x_error_smc = np.abs(errors_smc[:, 0])
    y_error_smc = np.abs(errors_smc[:, 1])
    z_error_smc = np.abs(errors_smc[:, 2])
    
    x_error_lqr = np.abs(errors_lqr[:, 0])
    y_error_lqr = np.abs(errors_lqr[:, 1])
    z_error_lqr = np.abs(errors_lqr[:, 2])
    
    # Steady-state error (last 2 seconds average)
    steady_start_idx = int(0.8 * len(error_norm_pid))
    
    # 1. 扰动抑制比 (Disturbance Rejection Ratio)
    def calculate_disturbance_rejection_ratio(error_norm, disturbance_type):
        """
        计算扰动抑制比
        在扰动期间的误差与无扰动期间误差的比值
        """
        if disturbance_type == 'none':
            return 1.0
            
        # 定义扰动期间
        disturbance_indices = []
        if disturbance_type == 'impulse':
            disturbance_indices = list(range(200, 210)) + list(range(500, 510))
        elif disturbance_type == 'step':
            disturbance_indices = list(range(300, len(error_norm)))
        elif disturbance_type == 'constant':
            disturbance_indices = list(range(0, len(error_norm)))
        else:
            # 对于其他扰动类型，使用误差较大的时期
            high_error_indices = np.where(error_norm > np.percentile(error_norm, 70))[0]
            disturbance_indices = high_error_indices
        
        if len(disturbance_indices) == 0:
            return 1.0
            
        # 无扰动期间的误差（前1秒）
        normal_error = np.mean(error_norm[:100])
        disturbance_error = np.mean(error_norm[disturbance_indices])
        
        if normal_error < 1e-6:
            return disturbance_error
            
        return disturbance_error / normal_error
    
    # 2. 误差积分指标 (Integral Error Metrics)
    def calculate_integral_metrics(error_norm):
        """计算各种误差积分指标"""
        # ITAE (时间乘绝对误差积分) - 惩罚持续误差
        itae = np.sum(np.arange(len(error_norm)) * 0.01 * error_norm)
        
        # ISE (误差平方积分) - 强调大误差
        ise = np.sum(error_norm**2)
        
        # IAE (绝对误差积分) - 总误差量
        iae = np.sum(error_norm)
        
        return itae, ise, iae
    
    # 3. 稳定时间 (Settling Time)
    def calculate_settling_time(error_norm, percentage=5):
        """计算稳定到最终值±percentage%的时间"""
        final_value = np.mean(error_norm[steady_start_idx:])
        if final_value < 1e-6:
            return 0.0
            
        threshold = final_value * (1 + percentage/100)
        
        # 找到最后一个超过阈值的时间点
        exceeding_indices = np.where(error_norm > threshold)[0]
        if len(exceeding_indices) == 0:
            return 0.0
            
        last_exceed_idx = exceeding_indices[-1]
        return last_exceed_idx * 0.01
    
    # 4. 鲁棒性裕度 (Robustness Margin)
    def calculate_robustness_margin(error_norm):
        """基于误差的波动程度计算鲁棒性裕度"""
        steady_state_error = error_norm[steady_start_idx:]
        if len(steady_state_error) == 0:
            return 0.0
            
        # 使用变异系数(标准差/均值)的倒数作为鲁棒性指标
        cv = np.std(steady_state_error) / (np.mean(steady_state_error) + 1e-6)
        return 1.0 / (1.0 + cv)  # 值越大表示鲁棒性越好
    
    # 计算各个指标
    disturbance_rejection_pid = calculate_disturbance_rejection_ratio(error_norm_pid, disturbance_type)
    disturbance_rejection_smc = calculate_disturbance_rejection_ratio(error_norm_smc, disturbance_type)
    disturbance_rejection_lqr = calculate_disturbance_rejection_ratio(error_norm_lqr, disturbance_type)
    
    itae_pid, ise_pid, iae_pid = calculate_integral_metrics(error_norm_pid)
    itae_smc, ise_smc, iae_smc = calculate_integral_metrics(error_norm_smc)
    itae_lqr, ise_lqr, iae_lqr = calculate_integral_metrics(error_norm_lqr)
    
    settling_time_pid = calculate_settling_time(error_norm_pid)
    settling_time_smc = calculate_settling_time(error_norm_smc)
    settling_time_lqr = calculate_settling_time(error_norm_lqr)
    
    robustness_margin_pid = calculate_robustness_margin(error_norm_pid)
    robustness_margin_smc = calculate_robustness_margin(error_norm_smc)
    robustness_margin_lqr = calculate_robustness_margin(error_norm_lqr)
    
    # 振荡指标
    oscillation_pid = np.std(error_norm_pid[steady_start_idx:])
    oscillation_smc = np.std(error_norm_smc[steady_start_idx:])
    oscillation_lqr = np.std(error_norm_lqr[steady_start_idx:])
    
    # 超调量
    steady_error_pid = np.mean(error_norm_pid[steady_start_idx:])
    steady_error_smc = np.mean(error_norm_smc[steady_start_idx:])
    steady_error_lqr = np.mean(error_norm_lqr[steady_start_idx:])
    
    overshoot_pid = (np.max(error_norm_pid) / steady_error_pid) if steady_error_pid > 0.001 else 0
    overshoot_smc = (np.max(error_norm_smc) / steady_error_smc) if steady_error_smc > 0.001 else 0
    overshoot_lqr = (np.max(error_norm_lqr) / steady_error_lqr) if steady_error_lqr > 0.001 else 0
    
    metrics = {
        'RMSE Position': [
            np.sqrt(np.mean(error_norm_pid**2)),
            np.sqrt(np.mean(error_norm_smc**2)),
            np.sqrt(np.mean(error_norm_lqr**2))
        ],
        'Max Position Error': [
            np.max(error_norm_pid),
            np.max(error_norm_smc),
            np.max(error_norm_lqr)
        ],
        'Steady-State Error': [
            steady_error_pid,
            steady_error_smc,
            steady_error_lqr
        ],
        'Disturbance Rejection Ratio': [
            disturbance_rejection_pid,
            disturbance_rejection_smc,
            disturbance_rejection_lqr
        ],
        'ITAE': [
            itae_pid,
            itae_smc,
            itae_lqr
        ],
        'ISE': [
            ise_pid,
            ise_smc,
            ise_lqr
        ],
        'Settling Time (s)': [
            settling_time_pid,
            settling_time_smc,
            settling_time_lqr
        ],
        'Robustness Margin': [
            robustness_margin_pid,
            robustness_margin_smc,
            robustness_margin_lqr
        ],
        'Oscillation Level': [
            oscillation_pid,
            oscillation_smc,
            oscillation_lqr
        ],
        'Overshoot Ratio': [
            overshoot_pid,
            overshoot_smc,
            overshoot_lqr
        ]
    }
    
    return metrics

def main():
    """Main function to run robustness tests with leg suspension"""
    disturbance_types = ['constant', 'sinusoidal', 'impulse', 'step', 'random', 'complex']
    trajectory_types = ['figure8', 'rough']
    
    print("WHEELED LEG ROBOT ROBUSTNESS TEST WITH LEG SUSPENSION")
    print("=" * 70)
    
    for traj_type in trajectory_types:
        print(f"\nTesting trajectory: {traj_type.upper()}")
        for dist_type in disturbance_types:
            try:
                metrics = run_robustness_test_with_legs(dist_type, traj_type)
                
                print(f"\nPerformance under {dist_type} disturbance:")
                print(f"{'Metric':<25} {'PID':<12} {'SMC':<12} {'LQR':<12}")
                print("-" * 65)
                for metric, values in metrics.items():
                    if metric in ['RMSE Position', 'Max Position Error', 'Steady-State Error']:
                        print(f"{metric:<25} {values[0]:<12.4f} {values[1]:<12.4f} {values[2]:<12.4f}")
            except Exception as e:
                print(f"Error in {dist_type} disturbance test: {e}")
                continue

if __name__ == "__main__":
    main()